#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <gmssl/sm2.h>
#include <gmssl/sm4.h>
#include <gmssl/sm3.h>
#include <gmssl/error.h>

int main(void) {

    // 1. 检查/生成 Alice SM2 密钥对
    FILE *fp_priv = fopen("alice_sm2.pem", "r");
    SM2_KEY alice_key;

    if (!fp_priv) {
        printf("Alice key not found — generating SM2 key pair...\n");
        sm2_key_generate(&alice_key);

        // 修正：GmSSL 3.1.1 只有三个参数
        FILE *f1 = fopen("alice_sm2.pem", "w");
        sm2_private_key_info_encrypt_to_pem(&alice_key, "1234", f1);
        fclose(f1);

        FILE *f2 = fopen("alice_sm2_pub.pem", "w");
        sm2_public_key_info_to_pem(&alice_key, f2);
        fclose(f2);

        printf("Alice SM2 keypair generated: alice_sm2.pem, alice_sm2_pub.pem\n");
    } else {
        sm2_private_key_info_decrypt_from_pem(&alice_key, "1234", fp_priv);
        fclose(fp_priv);
        printf("Alice SM2 keypair loaded.\n");
    }

    // 2. 读取 Bob 的公钥
    FILE *fp_bobpub = fopen("bob_sm2_pub.pem", "r");
    if (!fp_bobpub) {
        fprintf(stderr, "找不到 Bob 公钥文件 bob_sm2_pub.pem\n");
        return -1;
    }

    SM2_KEY bob_pub;
    sm2_public_key_info_from_pem(&bob_pub, fp_bobpub);
    fclose(fp_bobpub);

    // 3. 读取密文、数字信封、签名
    FILE *f_cipher = fopen("ciphertext.bin", "rb");
    FILE *f_ek = fopen("ek.bin", "rb");
    FILE *f_sig = fopen("sig.bin", "rb");

    if (!f_cipher || !f_ek || !f_sig) {
        fprintf(stderr, "缺少 Bob 生成的文件，请确保 ciphertext.bin、ek.bin、sig.bin 已生成。\n");
        return -1;
    }

    unsigned char ciphertext[1024], ek[512], sig[SM2_MAX_SIGNATURE_SIZE];
    size_t cipherlen = fread(ciphertext, 1, sizeof(ciphertext), f_cipher);
    size_t eklen = fread(ek, 1, sizeof(ek), f_ek);
    size_t siglen = fread(sig, 1, sizeof(sig), f_sig);
    fclose(f_cipher);
    fclose(f_ek);
    fclose(f_sig);

    // 4. 验签
    unsigned char dgst[32];
    sm3_digest(ciphertext, cipherlen, dgst);

    if (sm2_verify(&bob_pub, dgst, sig, siglen) == 1) {
        printf("签名验证成功！来自 Bob 的消息真实可信。\n");
    } else {
        fprintf(stderr, "签名验证失败！可能数据被篡改。\n");
        return -1;
    }

    // 5. 解密数字信封（取出 key + iv）
    unsigned char keyiv[32];
    size_t keyivlen = sizeof(keyiv);

    if (sm2_decrypt(&alice_key, ek, eklen, keyiv, &keyivlen) != 1) {
        fprintf(stderr, "SM2 解密数字信封失败\n");
        return -1;
    }

    unsigned char key[16], iv[16];
    memcpy(key, keyiv, 16);
    memcpy(iv, keyiv + 16, 16);

    // 6. SM4 解密
    unsigned char *plain = calloc(1, cipherlen);
    SM4_KEY sm4_dk;
    sm4_set_decrypt_key(&sm4_dk, key);
    sm4_cbc_decrypt(&sm4_dk, iv, ciphertext, cipherlen / 16, plain);

    FILE *fout = fopen("message_dec.txt", "w");
    fwrite(plain, 1, cipherlen, fout);
    fclose(fout);

    printf("Alice 已成功解密密文，结果保存至 message_dec.txt\n");

    free(plain);
    return 0;
}
